# %%
# -*- coding: utf-8 -*-
import numpy as np
import xlrd
import time
import os
import collections
from matplotlib import pyplot as plt
from scipy import ndimage as nd
import scipy.io as sio
from scipy.optimize import curve_fit
import socket
os.system('chcp 65001')
Host_Name = socket.gethostname()
if Host_Name == 'VictorChengPC':
    Computer_Name = 'X:\\Work\\mail.ustc.edu.cn\\Rb87BEC - Files\\'
elif Host_Name == 'VictorChengXPS':
    Computer_Name = 'X:\\Work\\mail.ustc.edu.cn\\Rb87BEC - Files\\'
elif Host_Name == 'calcenter2':
    Computer_Name = 'D:\\vc\\mail.ustc.edu.cn\\Rb87BEC - Files\\'
# %%
time_start = time.time()
time_start_struct = time.localtime(time_start)

# %%
Filter_Size = 2
# Computer_Name = '\\\\calcenter2\\'
Fit_Prefix = ['NoFit', '']
Fit_BEC_Switch = 0
Data_Table = xlrd.open_workbook(Computer_Name + 'Rb87files\\3DSOC_Quench\\Quench.xlsx').sheets()[0]
TimeJoin_Table = xlrd.open_workbook(Computer_Name + 'Rb87files\\3DSOC_Quench\\Quench.xlsx').sheets()[1]
Meta_Table = xlrd.open_workbook(Computer_Name + 'Rb87files\\3DSOC_Quench\\Quench.xlsx').sheets()[2]
SaveFigure_Address_Selected = '\\\\calcenter2\\Rb87figures\\3DSOC_Quench\\figures' + Fit_Prefix[Fit_BEC_Switch] + \
                              '\\Selected_Pic\\'
os.system('mkdir ' + '\"' + SaveFigure_Address_Selected + '\"')

SaveData_Address_Selected = Computer_Name + 'Rb87Data\\3DSOC_Quench\\' + Fit_Prefix[Fit_BEC_Switch] + \
                              '\\Selected\\'
os.system('mkdir ' + '\"' + SaveData_Address_Selected + '\"')

Meta_Row_List = range(27, 27 + 1)

# %%
for Meta_Row_Index in range(len(Meta_Row_List)):
    Meta_Row = Meta_Row_List[Meta_Row_Index]
    Meta = Meta_Table.row_values(Meta_Row - 1)
    Try = int(Meta[1])
    Detuning = Meta[2]
    Mat_Data_command = Meta[3]
    Mat_Data = []
    TimeJoinRow_List = []
    exec(Mat_Data_command)
    Bad_Dict_Command = Meta[4]
    Bad_Dict = {}
    exec(Bad_Dict_Command)
    Data_Source = Meta[6]
    Photo_AllTime_1d = collections.OrderedDict()
    Photo_rened_AllTime_1d = collections.OrderedDict()
# Photo
    if Data_Source == 'Photo':
        for Line_Index in range(len(Mat_Data)):
            Line = Mat_Data[Line_Index]
            Input_Matrix = Data_Table.row_values(Line - 1)
            YearMonthDay = str(int(Input_Matrix[0]))
            Folder_Name = Computer_Name + 'Rb87figures\\3DSOC_Quench\\figures' + \
                Fit_Prefix[Fit_BEC_Switch] + '\\' + YearMonthDay + '\\All_Spin_Polarization\\'
            Ini_Photo = int(Input_Matrix[1])
            Fin_Photo = int(Input_Matrix[2])
            V0 = Input_Matrix[4]
            Vy = Input_Matrix[5]
            Omega = Input_Matrix[6]
            Detuning = Input_Matrix[10]
            Cam = Input_Matrix[8]
            Time_Start = int(Input_Matrix[11])
            Time_Interval = int(Input_Matrix[12])
            Time_End = int(Input_Matrix[13])
            if Line in Bad_Dict.keys():
                Bad_List = Bad_Dict[Line]
            else:
                Bad_List = []
            for Photo_Count in range(Fin_Photo - Ini_Photo + 1):
                Photo_Ind = Photo_Count + Ini_Photo
                Time = (Photo_Ind - Ini_Photo) * Time_Interval + Time_Start
                if Time not in Photo_AllTime_1d.keys():
                    Photo_AllTime_1d[Time] = []
                if Photo_Ind in Bad_List:
                    continue
                Data = np.load(
                    Folder_Name + 'Cam_' + Cam + '_photo%d_Spin_Polarization_FBZCircleRectangle_1d.npz' % Photo_Ind)
                SP_1d = Data['SP_1d']
                Photo_AllTime_1d[Time].append(SP_1d)
        #            if Time in Photo_AllTime.keys():
        #                Photo_AllTime[Time].append(SP)
        #            else:
        #                Photo_AllTime[Time] = [SP]
        SP_AllTime_Ave_1d = collections.OrderedDict()
        for Time_Index in range(len(Photo_AllTime_1d.keys())):
            Time = list(Photo_AllTime_1d.keys())[Time_Index]
            SP_ThisTime_1d = np.array(Photo_AllTime_1d[Time])
            SP_ThisTime_Ave_1d = np.mean(SP_ThisTime_1d, 0)
            SP_AllTime_Ave_1d[Time] = SP_ThisTime_Ave_1d
# TimeJoin
    elif Data_Source == 'TimeJoin':
        Cam = Meta[5]
        Parameter_ref = int(Meta[7])
        Parameter_ref_Input_Matrix = Data_Table.row_values(Parameter_ref - 1)
        V0 = Parameter_ref_Input_Matrix[4]
        Vy = Parameter_ref_Input_Matrix[5]
        Omega = Parameter_ref_Input_Matrix[6]
        Detuning = Parameter_ref_Input_Matrix[10]
        TimeJoin_DataFolder = Computer_Name + 'Rb87Data\\3DSOC_Quench\\NoFit\\TimeJoin\\'
        for TimeJoinRow_Index in range(len(TimeJoinRow_List)):
            TimeJoinRow = TimeJoinRow_List[TimeJoinRow_Index]
            Data = np.load(TimeJoin_DataFolder + 'TimeJoinRow=%d_1d.npz' % TimeJoinRow, allow_pickle=True)
            SP_TimeJoin_1d = Data['SP_TimeJoin_1d'].item()
            SP_rened_TimeJoin_1d = Data['SP_rened_TimeJoin_1d'].item()
            ts_TimeJoin = Data['ts_TimeJoin']
            Photo_Ind_List = list(SP_TimeJoin_1d.keys())
            if TimeJoinRow in Bad_Dict:
                Bad_List = Bad_Dict[TimeJoinRow]
            else:
                Bad_List = []
            if 'Bad_Dict' in Data:
                Bad_List = np.concatenate(Bad_List, Data['Bad_List'])
            for t_Index in range(len(ts_TimeJoin)):
                Time = ts_TimeJoin[t_Index]
                if t_Index >= len(Photo_Ind_List):
                    continue
                Photo_Ind = Photo_Ind_List[t_Index]
                if Time not in Photo_AllTime_1d.keys():
                    Photo_AllTime_1d[Time] = []
                    Photo_rened_AllTime_1d[Time] = []
                if int(Photo_Ind) in Bad_List:
                    continue
                SP_1d = SP_TimeJoin_1d[Photo_Ind]
                SP_rened_1d = SP_rened_TimeJoin_1d[Photo_Ind]
                # SP_filtered_1d = nd.filters.generic_filter(SP_1d, np.nanmean, size=Filter_Size, mode='mirror')
                # SP_rened_1d = SP_filtered_1d[Filter_Size // 2:np.shape(SP_filtered_1d)[0]:Filter_Size]
                Photo_AllTime_1d[Time].append(SP_1d)
                Photo_rened_AllTime_1d[Time].append(SP_rened_1d)
        SP_AllTime_Ave_1d = collections.OrderedDict()
        SP_AllTime_Std_1d = collections.OrderedDict()
        SP_rened_AllTime_Ave_1d = collections.OrderedDict()
        SP_rened_AllTime_Std_1d = collections.OrderedDict()
        for Time in Photo_AllTime_1d:
            SP_ThisTime_1d = np.array(Photo_AllTime_1d[Time])
            SP_AllTime_Ave_1d[Time] = np.mean(SP_ThisTime_1d, 0)
            SP_AllTime_Std_1d[Time] = np.std(SP_ThisTime_1d, 0)

            SP_rened_ThisTime_1d = np.array(Photo_rened_AllTime_1d[Time])
            SP_rened_AllTime_Ave_1d[Time] = np.mean(SP_rened_ThisTime_1d, 0)
            SP_rened_AllTime_Std_1d[Time] = np.std(SP_rened_ThisTime_1d, 0)

        np.savez(SaveData_Address_Selected + 'Cam_' + Cam + '_' + 'Meta_Row=%d,Filter_Size=%s_1d' %
                 (Meta_Row, str(Filter_Size)),
                 SP_AllTime_Ave_1d=SP_AllTime_Ave_1d, SP_AllTime_Std_1d=SP_AllTime_Std_1d,
                 SP_rened_AllTime_Ave_1d=SP_rened_AllTime_Ave_1d,
                 SP_rened_AllTime_Std_1d=SP_rened_AllTime_Std_1d,
                 ts_TimeJoin=ts_TimeJoin)  # , Photo_rened_AllTime=Photo_rened_AllTime)
        sio.savemat(SaveData_Address_Selected + 'Cam_' + Cam + '_' + 'Meta_Row=%d,Filter_Size=%s.mat' %
                    (Meta_Row, str(Filter_Size)),
                    {'SP_AllTime_Ave_1d': np.array(list(SP_AllTime_Ave_1d.values())),
                     'SP_AllTime_Std_1d': np.array(list(SP_AllTime_Std_1d.values())),
                     'SP_rened_AllTime_Ave_1d': np.array(list(SP_rened_AllTime_Ave_1d.values())),
                     'SP_rened_AllTime_Std_1d': np.array(list(SP_rened_AllTime_Std_1d.values())),
                     'ts_TimeJoin': ts_TimeJoin})

# %%
time_end = time.time()
time_end_struct = time.localtime(time_end)
print('Time spent: ', time_end - time_start, 's')
print('Start time:', time_start_struct)
print('End time:', time_end_struct)
